'''
Created on Aug 31, 2009

@author: mkiyer
'''
import xml.etree.cElementTree as etree

def parse_elem_samplegroup(elem, sdbi):
    # parse the xml elements
    assert elem.tag == 'samplegroup'
    name_elem = elem.find('name')
    name = name_elem.text if name_elem is not None else None
    color_elem = elem.find('color')
    color = color_elem.text if color_elem is not None else None
    cond_elems = elem.findall('cond')
    conds = []
    if cond_elems is not None:
        for cond_elem in cond_elems:
            cond = {}
            for e in cond_elem.getchildren():
                cond[e.tag] = e.text
            conds.append(cond)
    samples = [x.text for x in elem.findall('sample')]
    # now build a samplegroup object from the elements
    return SampleGroup(sdbi, name, color, samples, conds)

def parse_samplegroups_xml(xmlfile, sdbi):
    tree = etree.parse(xmlfile)
    root = tree.getroot()
    assert root.tag == 'samplegroups'    
    for e in root.getchildren():
        yield parse_elem_samplegroup(e, sdbi)

def parse_elem_experiment(elem):
    assert elem.tag == 'experiment'
    name_elem = elem.find('name')
    name = name_elem.text if name_elem is not None else None
    exper_elem = elem.find('exper')
    exper = exper_elem.text if exper_elem is not None else None
    control_elem = elem.find('control')
    control = control_elem.text if control_elem is not None else None
    return name, exper, control

def parse_experiment_xml(xmlfile, sdbi):
    tree = etree.parse(xmlfile)
    root = tree.getroot()
    assert root.tag == 'experiments'
    sample_groups = {}
    for elem in root.findall('samplegroup'):
        sg = parse_elem_samplegroup(elem, sdbi) 
        sample_groups[sg.name] = sg
    for elem in root.findall('experiment'):
        name, exper, control = parse_elem_experiment(elem)
        assert exper in sample_groups
        assert control in sample_groups
        yield Experiment(name, sample_groups[exper], sample_groups[control])

class Experiment(object):
    def __init__(self, name=None, exper=None, control=None):
        self.name = name
        self.exper = exper
        self.control = control

class SampleGroup(object):
    '''
    Represents a group of samples and associates each group with a name
    and color for plotting
    
    A SampleGroup is specified by either a list of sample names or a 
    list of conditions (conds) that use SampleDB parameters to describe a
    class of samples
    '''
    def __init__(self, sdbi, name=None, color='black', samples=None, conds=None):
        self.name = name
        self.color = color
        self.samples = set(sdbi.get_sample_names(samples, conds))

if __name__ == '__main__':
    import sys
    groups = [x for x in parse_samplegroups_xml(sys.argv[1])]
    